package com.demo.java.limiter;

import com.demo.java.utils.LogUtils;
import com.demo.java.utils.ThreadUtils;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 漏斗限流
 * 漏斗容器固定，漏斗漏出速度固定，漏出后判断漏斗剩余容量是大于请求数，大于请求数则放行，否则请求失败
 *
 * 问题：漏斗漏出速度固定，不能应对后端服务能力的提升，如后端处理能力由1000QPS提升到 10000QPS，漏斗同不能解决
 *
 */
public class LeakBucketLimiterTest {

    public static void main(String[] args) throws InterruptedException {
        LeakBucketLimiter limiter = new LeakBucketLimiter();
        ExecutorService executor = Executors.newFixedThreadPool(10);
        int threads = 2;
        int loop = 20;
        AtomicInteger limited = new AtomicInteger(0);
        long start = System.currentTimeMillis();

        CountDownLatch latch = new CountDownLatch(threads);


        for (int i = 0; i < threads; i++) {
            executor.submit(() -> {
                for (int j = 0; j < loop; j++) {
                    if (!limiter.tryAcquire(1)) {
                        limited.incrementAndGet();
                    }
                    ThreadUtils.sleep(200, TimeUnit.MILLISECONDS);
                }
                System.out.println(LogUtils.logt("执行结束"));
                latch.countDown();
            });
        }


        System.out.println(LogUtils.log("等待结束"));
        latch.await();

        long cost = (System.currentTimeMillis() - start)/1000;
        System.out.println(LogUtils.log("被限制次数: " + limited.get()) + ", 通过的次数: " + (threads * loop - limited.get()));
        System.out.println(LogUtils.log("限制比例: " + (limited.get()/(float)(threads * loop))));
        System.out.println(LogUtils.log("花费时间: " + cost));

        executor.shutdown();
    }

    static class LeakBucketLimiter {
        // 漏斗漏出时间
        private long leakOut = System.currentTimeMillis();
        // 漏斗漏出速度 /s
        private int leakRate = 2;
        // 漏斗容量
        private int capacity = 2;
        // 漏斗中剩余水的数量
        private AtomicInteger water = new AtomicInteger();

        public boolean tryAcquire(int count) {
            long now = System.currentTimeMillis();
            // 空桶直接获取成功
            if (water.get() == 0) {
                leakOut = now;
                water.incrementAndGet();
                return true;
            }
            /** 漏斗漏出 */
            // 计算单位时间漏出个数
            int leakCount = (int)(((now - leakOut) / 1000) * leakRate);
            int waterLeak = water.get() - leakCount;
            // 漏斗剩余数量整形，最小为0，不可为负数
            water.set(Math.max(0, waterLeak));
            // 更新漏出时间
            leakOut = now;

            // 判断剩余水量是否够用
            if (water.get() < capacity) {
                water.incrementAndGet();
                return true;
            }
            return false;
        }

    }
}
